from ship_envs.core.obs_viewer import ObsViewer
from ship_envs.core.ship import Ship
from ship_envs.environment.ImazuProblemEnv import ImazuProblemEnv
from ship_envs.environment.imazu_problem import case1, case2, case3, case4, case5, case6, case8, case7, case9, case10
from ship_envs.environment.imazu_problem_env_factory import new_case_env_by_index


def env_test(env):
    for _ in range(3):
        env.reset()
        while True:
            obs, reward, done, info = env.step(env.action_space.sample())
            print(obs)
            env.render()
            if done:
                break


def show_env_obs(env):
    env.reset()
    for _ in range(1):
        env.step(0)
    viewer = ObsViewer(500, 500)
    entitys = env.world.get_entitys()
    entitys.reverse()
    for entity in entitys:
        viewer.add_model(entity.model)
        if isinstance(entity, Ship):
            viewer.add_model(entity.target)
    img = viewer.get_img()
    img.show()


def case1_test():
    env = ImazuProblemEnv(case1())
    for _ in range(3):
        env.reset()
        while True:
            obs, reward, done, info = env.step(env.action_space.sample())
            env.render()
            if done:
                break


def show_case1_obs():
    env = ImazuProblemEnv(case1())
    env.reset()
    for _ in range(150):
        env.step(0)
    viewer = ObsViewer(500, 500)
    entitys = env.world.get_entitys()
    for entity in entitys:
        viewer.add_model(entity.model)
        if isinstance(entity, Ship):
            viewer.add_model(entity.target)
    img = viewer.get_img()
    img.show()

def reward_test(env):
    env.set_show_fairway(True)
    for _ in range(3):
        env.reset()
        while True:
            obs, reward, done, info = env.step(env.action_space.sample())
            print(obs, reward, done, info)
            env.render()
            if done:
                break


if __name__ == '__main__':
    # env = ImazuProblemEnv(case1())
    # env = ImazuProblemEnv(case2())
    # env = ImazuProblemEnv(case3())
    # env = ImazuProblemEnv(case4())
    # env = ImazuProblemEnv(case5())
    # env = ImazuProblemEnv(case6())
    # env = ImazuProblemEnv(case7())
    # env = ImazuProblemEnv(case8())
    # env = ImazuProblemEnv(case9())
    # env = ImazuProblemEnv(case10())
    # env_test(env)
    # show_env_obs(env)
    # case1_test()
    # show_case1_obs()
    env = new_case_env_by_index(1)
    # reward_test(env)
    env.obs_type = 1
    env_test(env)
